#include "curdfimporter.h"

#include <buola/scene/ctransform.h>
#include <buola/scene/cgeode.h>
#include <buola/scene/transform/crttransform.h>
#include <buola/scene/transform/crevolutionjoint.h>
#include <buola/scene/transform/cprismaticjoint.h>
#include <buola/scene/geodes/ccylinder.h>
#include <buola/xml.h>
#include <buola/zz.h>
#include <buola/zz/repeat.h>
#include <buola/zz/zreal.h>
#include <buola/zz/zchar.h>

namespace buola { namespace scene { namespace import {

CURDFImporter::CURDFImporter()
{

}

CURDFImporter::~CURDFImporter()
{
}

static mat::CVec3d parse_vec(xml::CNode pNode,const std::string &pAttrib,const mat::CVec3d &pDefault=mat::CVec3d(0,0,0))
{
    std::string lValue=pNode.GetStringAttrib(pAttrib);
    if(lValue.empty()) return pDefault;
    
    std::array<double,3> lValues;
    if(!zz::phrase_parse(lValue.begin(),lValue.end(),zz::repeat<3>(zz::dbl()),zz::space(),lValues))
    {
        msg_warn() << "URDF: can't parse vector\n";
        return pDefault;
    }
    
    return mat::CVec3d(lValues[0],lValues[1],lValues[2]);
}

static mat::CQuat_d parse_rot(xml::CNode pNode,const std::string &pAttrib)
{
    mat::CVec3d lVec=parse_vec(pNode,pAttrib);
    return mat::CQuat_d(mat::C3DRotation(lVec(2),lVec(1),lVec(0)));
    return mat::CQuat_d();
}

static PGeode parse_visual(xml::CNode pNode,const PRTTransform &pLink)
{
    PGeode lGeode=nullptr;
    
    for(auto lChild : pNode.Children())
    {
        if(lChild.MatchesName("origin"))
        {
            pLink->SetRotation(parse_rot(lChild,"rpy"));
            pLink->SetTranslation(parse_vec(lChild,"xyz"));
        }
        else if(lChild.MatchesName("geometry"))
        {
            for(auto lSub : lChild.Children())
            {
                if(lSub.MatchesName("mesh"))
                {
                    mat::CVec3d lScale=parse_vec(lSub,"scale",mat::CVec3d(1,1,1));
                    std::string lFilename=lSub.GetStringAttrib("filename");
                    if(lFilename.empty()) continue;
                    lGeode=CGeode::Import(io::file(lFilename),lScale(0));
                }
                else if(lSub.MatchesName("cylinder"))
                {
//                    lGeode=new CCylinder(lSub.GetDoubleAttrib("length"),lSub.GetDoubleAttrib("radius"),30,true);
                }
            }
        }
    }
    
    return lGeode;
}

struct CURDFLink
{
    PRTTransform mLink;
    PTransform mParent;
};

typedef std::map<std::string,CURDFLink> TLinkMap;

static void parse_link(xml::CNode pNode,TLinkMap &pMap)
{
    std::string lName=pNode.GetStringAttrib("name");
    if(lName.empty())
    {
        msg_warn() << "URDF: link node without name\n";
        return;
    }

    CURDFLink lLink;
    lLink.mLink=new CRTTransform;
    lLink.mLink->SetName(lName);
    
    for(auto lChild : pNode.Children())
    {
        if(lChild.MatchesName("visual"))
        {
            PGeode lGeode=parse_visual(lChild,lLink.mLink);
            if(lGeode)
                lGeode->AttachTo(lLink.mLink);
        }
    }

    pMap[lName]=lLink;
}

struct CURDFJoint
{
    std::string mParent;
    std::string mChild;
    PTransform mJoint;
};

typedef std::list<CURDFJoint> TJointList;

static void parse_joint(xml::CNode pNode,TJointList &pJoints)
{
    std::string lName=pNode.GetStringAttrib("name");
    std::string lType=pNode.GetStringAttrib("type");
    if(lName.empty()||lType.empty())
    {
        msg_warn() << "URDF: link node without name or type\n";
        return;
    }
    
    mat::CVec3d lOrigin(0,0,0);
    mat::CQuat_d lRot;
    mat::CVec3d lAxis(1,0,0);
    
    CURDFJoint lJoint;
    for(auto lChild : pNode.Children())
    {
        if(lChild.MatchesName("parent"))
        {
            lJoint.mParent=lChild.GetStringAttrib("link");
        }
        else if(lChild.MatchesName("child"))
        {
            lJoint.mChild=lChild.GetStringAttrib("link");
        }
        else if(lChild.MatchesName("origin"))
        {
            lOrigin=parse_vec(lChild,"xyz");
            lRot=parse_rot(lChild,"rpy");
        }
        else if(lChild.MatchesName("axis"))
        {
            lAxis=parse_vec(lChild,"xyz",mat::CVec3d(1,0,0));
        }
    }
    
    if(lType=="fixed")
    {
        lJoint.mJoint=new CRTTransform(lRot,lOrigin,lName);
    }
    else if(lType=="revolute"||lType=="continuous")
    {
        lJoint.mJoint=new CRevolutionJoint(lOrigin,lAxis,0,0,lName);
    }
    else if(lType=="prismatic")
    {
        lJoint.mJoint=new CPrismaticJoint(lOrigin,lAxis,0,0,lName);
    }
    else
    {
        msg_warn() << "URDF: unknown joint type " << lType << "\n";
        return;
    }
    
    pJoints.push_back(lJoint);
}

static PTransform parse_robot(xml::CNode pNode)
{
    if(!pNode.MatchesName("robot"))
        throw XData("error parsing URDF file: top level node is not robot");
    
    PRTTransform lRoot=new CRTTransform;

    TLinkMap lLinks;
    TJointList lJoints;
    
    for(auto lChild : pNode.Children())
    {
        if(lChild.Type()!=xml::ENodeType::ELEMENT) continue;
        
        if(lChild.MatchesName("link"))
        {
            parse_link(lChild,lLinks);
        }
        else if(lChild.MatchesName("joint"))
        {
            parse_joint(lChild,lJoints);
        }
    }

    //first set all parents of links
    for(const auto &j : lJoints)
    {
        auto lChild=lLinks.find(j.mChild);
        
        if(lChild==lLinks.end())
        {
            msg_warn() << "URDF: child link not found in joint description\n";
            continue;
        }
        
        if(lChild->second.mParent)
            msg_warn() << "URDF: link with two parents\n";
        j.mJoint->AddChild(lChild->second.mLink);
        lChild->second.mParent=j.mJoint;
    }
    
    //links without a parent are set to be children of root
    for(auto &l : lLinks)
    {
        if(l.second.mParent) continue;
        l.second.mParent=lRoot;
        msg_info() << l.first << "\n";
        lRoot->AddChild(l.second.mLink);
    }
    
    //now set all joints as the child of another joint
    for(const auto &j : lJoints)
    {
        auto lParent=lLinks.find(j.mParent);
        
        if(lParent==lLinks.end())
        {
            msg_warn() << "URDF: link not found in joint description\n";
            continue;
        }
        
        lParent->second.mParent->AddChild(j.mJoint);
    }
    
    
    return lRoot;
}

PTransform CURDFImporter::Parse(const io::CURI &pURI)
{
    xml::CDoc lXML;
    lXML.Load(pURI);

    return parse_robot(lXML.Root());
}

/*namespace import*/ } /*namespace scene*/ } /*namespace buola*/ }
